# 导入所需的库
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import yfinance as yf


# 定义一个函数来计算唐奇安通道
def calcDonchianChannels(data: pd.DataFrame, period: int):
    data["upperDon"] = data["High"].rolling(period).max()
    data["lowerDon"] = data["Low"].rolling(period).min()
    data["midDon"] = (data["upperDon"] + data["lowerDon"]) / 2
    return data


# 获取一些股票数据
ticker = input("请输入股票编号,大A用'.SS'结尾: ")
yfObj = yf.Ticker(ticker)
data = yfObj.history().drop(
    ["Volume", "Stock Splits"], axis=1)
print(f'{data}')
# 计算20日唐奇安通道
data = calcDonchianChannels(data, 10)
print("开始计算===========")

# 定义一个函数来生成交易信号和回报率
def donchianExitStrategy(data: pd.DataFrame, exit_days: int):
    # 初始化一个列来存储交易信号和持仓状态
    data["signal"] = 0  # 1表示多头，-1表示空头，0表示无仓位
    data["position"] = 0  # 同上

    # 初始化一个列来存储入场时间和价格
    data["entry_time"] = None
    data["entry_price"] = None

    # 遍历每一行数据，判断交易条件并执行交易逻辑
    for i in range(len(data)):
        # 获取当前行的数据
        row = data.iloc[i]

        # 如果当前没有仓位，检查是否有入场信号（价格突破中间值）
        if data.iloc[i - 1]["position"] == 0:
            if row["Close"] > row["midDon"]:
                # 多头入场信号，买入并记录入场时间和价格
                data.loc[row.name, "signal"] = 1
                data.loc[row.name, "position"] = 1
                data.loc[row.name, "entry_time"] = row.name
                data.loc[row.name, "entry_price"] = row["Close"]
            elif row["Close"] < row["midDon"]:
                # 空头入场信号，卖出并记录入场时间和价格
                data.loc[row.name, "signal"] = -1
                data.loc[row.name, "position"] = -1
                data.loc[row.name, "entry_time"] = row.name
                data.loc[row.name, "entry_price"] = row["Close"]

        # 如果当前有多头仓位，检查是否有出场信号（到达指定天数或者价格跌破下界）
        elif data.iloc[i - 1]["position"] == 1:
            if (row.name - pd.to_datetime(data.iloc[i - 1]["entry_time"])) >= pd.Timedelta(days=exit_days) or \
                    row["Close"] < row["lowerDon"]:
                # 多头出场信号，卖出并清空入场时间和价格记录
                data.loc[row.name, "signal"] = -1
                data.loc[row.name, "position"] = 0
                data.loc[row.name, "entry_time"] = None
                data.loc[row.name, "entry_price"] = None

        # 如果当前有空头仓位，检查是否有出场信号（到达指定天数或者价格突破上界）
        elif data.iloc[i - 1]["position"] == -1:
            if (row.name - pd.to_datetime(data.iloc[i - 1]["entry_time"])) >= pd.Timedelta(days=exit_days) or \
                    row["Close"] > row["upperDon"]:
                # 空头出场信号，买入并清空入场时间和价格记录
                data.loc[row.name, "signal"] = 1
                data.loc[row.name, "position"] = 0
                data.loc[row.name, "entry_time"] = None
                data.loc[row.name, "entry_price"] = None

        #    如果当前没有交易信号，沿用上一行的持仓状态
            else:
                data.loc[row.name, "position"] = data.iloc[i - 1]["position"]

        # 计算每日回报率和累积回报率
        data["returns"] = np.log(data["Close"] / data["Close"].shift(1))
        data["strategy_returns"] = data["returns"] * data["position"].shift(1)
        data["cum_returns"] = np.exp(data["returns"].cumsum()) - 1
        data["cum_strategy_returns"] = np.exp(data["strategy_returns"].cumsum()) - 1

        return data

# 使用80天作为出场条件，调用策略函数，并打印最后几行数据
data_80d_exit = donchianExitStrategy(data.copy(), exit_days=80)
# 假设你已经有了一个数据框data_80d_exit，包含了唐奇安通道和交易信号
# 你可以用以下的代码来加入买入价的区间
print(f"最后值为:{data_80d_exit.tail()}")

# 定义一个参数来表示区间的宽度，比如0.1表示10%
width = 0.1

# 计算区间的上下限
data_80d_exit["upper_limit"] = data_80d_exit["midDon"] + width * (data_80d_exit["upperDon"] - data_80d_exit["lowerDon"])
data_80d_exit["lower_limit"] = data_80d_exit["midDon"] - width * (data_80d_exit["upperDon"] - data_80d_exit["lowerDon"])

# 根据区间来修改交易信号
data_80d_exit["signal"] = np.where((data_80d_exit["Close"] > data_80d_exit["upper_limit"]) & (data_80d_exit["position"].shift(1) != 1), 1, data_80d_exit["signal"])
data_80d_exit["signal"] = np.where((data_80d_exit["Close"] < data_80d_exit["lower_limit"]) & (data_80d_exit["position"].shift(1) != -1), -1, data_80d_exit["signal"])
# 找到所有多头或者空头的买入信号（signal为1）
buy_signals = data_80d_exit[data_80d_exit["signal"] == 1]
# 如果有买入信号，取最后一行的收盘价作为最近的买入价格
if len(buy_signals) > 0:
    last_buy_price = buy_signals.iloc[-1]["Close"]
    print(f"最近一次买入价格为：{last_buy_price:.2f}")
else:
    print("没有找到任何买入信号")
print(f"最后值为:{data_80d_exit.tail()}")


